import networkx as nx
import matplotlib.pyplot as plt
import numpy as np
import pulp as pl
import random
import math
from collections import defaultdict
from matplotlib.font_manager import FontProperties
font = FontProperties(fname='resources/SimHei.ttf', size=12)  # 替换为实际字体路径

# 1. 生成有向管网拓扑（确保流向高压→低压）
def generate_gas_network(num_nodes=50, num_sources=3, num_users=15, valve_prob=0.4):
    """生成带压力梯度的有向管网"""
    # 创建空的有向图
    G = nx.DiGraph()
    
    # 生成节点位置（确保气源在中心区域）
    center = (0.5, 0.5)
    positions = {}
    for i in range(num_nodes):
        # 气源节点集中在中心，其他节点随机分布
        if i < num_sources:
            positions[i] = (
                center[0] + random.uniform(-0.2, 0.2),
                center[1] + random.uniform(-0.2, 0.2)
            )
        else:
            positions[i] = (random.random(), random.random())
    
    # 设置节点压力梯度（气源 > 普通节点 > 用户）
    for i in range(num_nodes):
        if i < num_sources:
            G.add_node(i, type='source', pressure=random.uniform(0.8, 1.2), pos=positions[i])
        elif i < num_sources + num_users:
            G.add_node(i, type='user', pressure=random.uniform(0.4, 0.6), 
                      priority=random.randint(1, 5), pressure_req=random.uniform(0.2, 0.4),
                      pos=positions[i])
        else:
            G.add_node(i, type='junction', pressure=random.uniform(0.5, 0.7), pos=positions[i])
    
    # 添加带流向的管道（高压→低压）
    for i in range(num_nodes):
        for j in range(i+1, num_nodes):
            if random.random() < 0.3:  # 连接概率
                dist = math.dist(positions[i], positions[j])
                if dist < 0.4:  # 只连接邻近节点
                    # 根据压力确定流向
                    if G.nodes[i]['pressure'] > G.nodes[j]['pressure']:
                        u, v = i, j
                    else:
                        u, v = j, i
                    
                    # 添加管道属性
                    diameter = random.choice([0.1, 0.15, 0.2, 0.25])
                    length = dist * 10  # 实际距离(km)
                    resistance = random.uniform(0.02, 0.05)
                    
                    # 添加带方向的边
                    G.add_edge(u, v, 
                              diameter=diameter,
                              length=length,
                              resistance=resistance,
                              valve=random.random() < valve_prob)
                    
                    # 设置阀门属性
                    if G.edges[u, v]['valve']:
                        G.edges[u, v]['valve_time'] = random.choice([1, 2, 3, 5, 10])
                        G.edges[u, v]['operational_cost'] = random.uniform(0.5, 2.0)
    
    return G, positions

# 2. 流向验证函数（防止逆流）
def validate_flow_direction(G, sources):
    """确保所有路径都是从气源向外辐射"""
    for node in G.nodes:
        # 跳过气源点
        if node in sources:
            continue
            
        # 检查是否存在从气源到该节点的路径
        has_path_from_source = any(nx.has_path(G, source, node) for source in sources)
        
        # 检查是否存在指向气源的路径（禁止逆流）
        has_path_to_source = any(nx.has_path(G, node, source) for source in sources)
        
        if has_path_to_source and not has_path_from_source:
            # 修复逆流：断开指向气源的边
            for source in sources:
                if G.has_edge(node, source):
                    G.remove_edge(node, source)
                    print(f"Removed reverse flow edge: {node}->{source}")
    
    # 二次验证
    for source in sources:
        descendants = nx.descendants(G, source)
        ancestors = nx.ancestors(G, source)
        if ancestors:
            print(f"Warning: Source {source} has incoming paths from {ancestors}")
    return G

# 3. 泄漏点模拟
def identify_leaks(G, num_leaks=2):
    """在非气源节点生成泄漏点"""
    non_source_nodes = [n for n in G.nodes if G.nodes[n]['type'] != 'source']
    leak_nodes = random.sample(non_source_nodes, num_leaks)
    
    leaks = {}
    for node in leak_nodes:
        pressure = G.nodes[node]['pressure']
        leak_rate = random.uniform(0.8, 1.2) * pressure * math.sqrt(pressure)
        leaks[node] = (G.nodes[node]['pos'], leak_rate)
    
    return leaks

# 4. 最小割算法（有向图版本）
def find_min_cut(G, leaks, sources):
    """基于流向的有向图最小割算法"""
    flow_graph = nx.DiGraph()
    super_source = "super_source"
    super_sink = "super_sink"
    flow_graph.add_node(super_source)
    flow_graph.add_node(super_sink)
    
    # 构建带容量的流向图
    for u, v in G.edges():
        capacity = 1 if G.edges[u, v].get('valve', False) else 10000
        flow_graph.add_edge(u, v, capacity=capacity)
    
    # 连接超级源点->气源
    for source in sources:
        flow_graph.add_edge(super_source, source, capacity=10000)
    
    # 连接泄漏点->超级汇点
    for leak in leaks:
        flow_graph.add_edge(leak, super_sink, capacity=10000)
    
    # 计算最小割
    try:
        min_cut_value, partition = nx.minimum_cut(flow_graph, super_source, super_sink)
        reachable, non_reachable = partition
        
        # 提取割边（实际阀门）
        min_cut_set = []
        for u, v in flow_graph.edges():
            if u in reachable and v in non_reachable:
                if G.has_edge(u, v) and G.edges[u, v].get('valve', False):
                    min_cut_set.append((u, v))
        
        # 验证隔离区域
        H = G.copy()
        H.remove_edges_from(min_cut_set)
        isolated_nodes = set()
        for leak in leaks:
            if leak in H.nodes:
                isolated_nodes |= set(nx.descendants(H, leak)) | {leak}
                
        return min_cut_set, isolated_nodes
        
    except nx.NetworkXUnbounded:
        print("Error: Unbounded flow detected, check flow directions")
        return [], set()

# 5. 优化关阀策略（修复LpAffineExpression错误）
def optimize_valve_strategy(G, min_cut_set, leaks, sources):
    """多目标关阀优化（修复表达式错误）"""
    problem = pl.LpProblem("Gas_Valve_Optimization", pl.LpMinimize)
    
    # 决策变量
    valve_vars = {(u, v): pl.LpVariable(f"valve_{u}_{v}", 0, 1, pl.LpBinary)
                  for (u, v) in min_cut_set}
    
    # 添加最大关阀时间辅助变量
    t_max = pl.LpVariable("t_max", lowBound=0, cat='Continuous')
    
    # 目标函数组件
    user_penalty = 10
    leak_penalty = 8
    time_penalty = 1
    
    # 用户影响成本
    user_affect_cost = pl.lpSum(
        G.nodes[user]['priority'] for user in G.nodes 
        if G.nodes[user]['type'] == 'user'
    )
    
    # 泄漏成本（使用辅助变量）
    total_leak_rate = sum(leak_data[1] for leak_data in leaks.values())
    leak_cost = total_leak_rate * t_max
    
    # 时间成本
    time_cost = pl.lpSum(
        G.edges[edge]['valve_time'] * valve_vars[edge]
        for edge in min_cut_set
    )
    
    # 添加t_max约束
    for edge in min_cut_set:
        valve_time = G.edges[edge]['valve_time']
        problem += t_max >= valve_time * valve_vars[edge]
    
    # 目标函数
    problem += user_penalty * user_affect_cost + leak_penalty * leak_cost + time_penalty * time_cost
    
    # 约束：至少关闭一个阀门
    problem += pl.lpSum(valve_vars.values()) >= 1
    
    # 求解
    solver = pl.PULP_CBC_CMD(msg=False)
    problem.solve(solver)
    
    # 提取结果
    valves_to_close = [edge for edge, var in valve_vars.items() if var.value() == 1]
    
    # 计算受影响用户
    H = G.copy()
    H.remove_edges_from(valves_to_close)
    disconnected_users = []
    for user in [n for n, data in G.nodes(data=True) if data['type'] == 'user']:
        connected = any(nx.has_path(H, source, user) for source in sources)
        if not connected:
            disconnected_users.append(user)
            
    return valves_to_close, disconnected_users

# 6. 流向可视化
def visualize_network(G, pos, valves_to_close=None, leaks=None, disconnected_users=None, sources=None):
    """可视化管网与流向"""
    plt.figure(figsize=(15, 12))
    
    # 节点颜色分类
    node_colors = []
    for node in G.nodes:
        if leaks and node in leaks:
            node_colors.append('red')  # 泄漏点
        elif disconnected_users and node in disconnected_users:
            node_colors.append('purple')  # 断气用户
        elif sources and node in sources:
            node_colors.append('green')  # 气源
        elif G.nodes[node]['type'] == 'user':
            node_colors.append('blue')  # 正常用户
        else:
            node_colors.append('lightgray')  # 其他节点
    
    # 绘制节点
    nx.draw_networkx_nodes(G, pos, node_color=node_colors, node_size=300, edgecolors='black')
    
    # 绘制正常边（带箭头表示流向）
    nx.draw_networkx_edges(
        G, pos, 
        edgelist=[e for e in G.edges() if e not in (valves_to_close or [])],
        edge_color='gray',
        width=1.5,
        arrows=True,
        arrowstyle='->',
        arrowsize=15
    )
    
    # 绘制需关闭的阀门
    if valves_to_close:
        nx.draw_networkx_edges(
            G, pos, 
            edgelist=valves_to_close,
            edge_color='red',
            width=3,
            style='dashed',
            arrows=True,
            arrowstyle='->',
            arrowsize=20
        )
        # 标记阀门位置
        for u, v in valves_to_close:
            if u in pos and v in pos:
                mid_x = (pos[u][0] + pos[v][0]) / 2
                mid_y = (pos[u][1] + pos[v][1]) / 2
                plt.scatter(mid_x, mid_y, s=300, c='red', marker='X', zorder=3)
    
    # 压力梯度标注
    for node, (x, y) in pos.items():
        plt.text(x, y+0.02, f"{G.nodes[node]['pressure']:.2f}MPa", 
                fontsize=8, ha='center', color='darkgreen')
    
    # 图例
    plt.legend(handles=[
        plt.Line2D([0], [0], marker='o', color='w', markerfacecolor='green', markersize=10, label='气源'),
        plt.Line2D([0], [0], marker='o', color='w', markerfacecolor='blue', markersize=10, label='用户'),
        plt.Line2D([0], [0], marker='o', color='w', markerfacecolor='red', markersize=10, label='泄漏点'),
        plt.Line2D([0], [0], marker='o', color='w', markerfacecolor='purple', markersize=10, label='断气用户'),
        plt.Line2D([0], [0], color='gray', lw=2, label='正常管道'),
        plt.Line2D([0], [0], color='red', lw=2, linestyle='dashed', label='需关阀门')
    ], loc='upper right', fontsize=12, prop=font)
    
    plt.title("Gas Network Isolation Optimization", fontsize=14)
    plt.axis('off')
    plt.tight_layout()
    plt.savefig("gas_network_optimization.png", dpi=300)
    plt.show()

# 7. 主流程
def main():
    print("=== 燃气泄漏应急关阀优化系统 ===")
    
    print("1. 生成带压力梯度的有向管网...")
    G, pos = generate_gas_network(num_nodes=40, num_sources=3, num_users=12)
    
    sources = [n for n, data in G.nodes(data=True) if data['type'] == 'source']
    print("2. 验证流向（防止逆流）...")
    G = validate_flow_direction(G, sources)
    
    print("3. 模拟泄漏点...")
    leaks = identify_leaks(G, num_leaks=2)
    leak_nodes = list(leaks.keys())
    
    print("4. 查找最小阀门集合隔离泄漏点...")
    min_cut_set, isolated_nodes = find_min_cut(G, leak_nodes, sources)
    print(f"  初始需关闭的阀门数: {len(min_cut_set)}")
    
    print("5. 优化关阀策略...")
    valves_to_close, disconnected_users = optimize_valve_strategy(
        G, min_cut_set, leaks, sources
    )
    print(f"  优化后关闭的阀门数: {len(valves_to_close)}")
    print(f"  受影响用户数: {len(disconnected_users)}")
    
    print("6. 可视化优化结果...")
    visualize_network(G, pos, valves_to_close, leak_nodes, disconnected_users, sources)
    
    print("=== 优化完成 ===")

if __name__ == "__main__":
    main()